import math
import lpips
import numpy as np
import torch as pt
from skimage.metrics import structural_similarity, peak_signal_noise_ratio

def cumprod_exclusive(x):
  cp = pt.cumprod(x, 0)
  cp = pt.roll(cp, 1, 0)
  cp[0] = 1.0
  return cp

def random_sample(shape, ray_num):
  #shape [h, w]
  return pt.randperm(shape[0] * shape[1])[:ray_num]

def L_sample(shape, ray_num):
  #shape [h, w]
  # o o
  # o
  picks = pt.randperm((shape[0]-1) * (shape[1]-1))[:ray_num // 3]
  shift = picks // (shape[1]-1)
  picks += shift

  samples = pt.zeros(3 * len(picks)).to(pt.long)
  samples[0::3] = picks
  samples[1::3] = picks + 1
  samples[2::3] = picks + shape[1]
  return samples

@pt.no_grad()
def patch_render(render_func, rays_targ, R, T, sample_ray_N):
  ''' Render image in patch style.

  Args:
    rays_targ: [..., N, 3]; rays direction in target coordinate
    R: relative rotation, [..., 3, 3] matrices
    T: [..., 3], translations from target to source camera
      point p from target to source is accomplished via rot * p + t
  Returns:
    [..., N, C] images after bilinear sampling from input.
      Coordinates outside the image are sampled as 0
  '''
  predicted_image = []
  num = int(math.ceil(rays_targ.shape[-2] / sample_ray_N))
  for i in range(num):
    if i < num -1 :
      samples = pt.arange(sample_ray_N * i, sample_ray_N * (i + 1))
    else:
      samples = pt.arange(sample_ray_N * i, rays_targ.shape[-2])
    output = render_func(rays_targ[...,samples,:], R, T)
    predicted_image.append(output)
  predicted_image = pt.cat(predicted_image, 2)
  return predicted_image

def get_tvc_loss(rgbs):
  def totalVariation(images):
    pixel_dif1 = images[..., 1:, :, :] - images[..., :-1, :, :]
    pixel_dif2 = images[..., :, 1:, :] - images[..., :, :-1, :]
    sum_axis = [-3, -2, -1]

    tot_var = (
      pt.sum(pt.abs(pixel_dif1), dim=sum_axis) +
      pt.sum(pt.abs(pixel_dif2), dim=sum_axis))

    return tot_var / (images.shape[-2]-1) / (images.shape[-3]-1)

  return pt.mean(totalVariation(pt.sigmoid(rgbs)))

def measurement(rendered_image, gt_image):
  """ calculate PSNR/SSIM and LPIPS"""

  # calculate LPIPS. It's require to change image range from [0,1] to [-1,1]
  lpips_model = lpips.LPIPS(net='vgg')
  gt_lpips = gt_image.clone().cpu().permute(2, 0, 1) * 2.0 - 1.0
  predict_image_lpips = pt.Tensor(rendered_image).permute(2, 0, 1) * 2.0 - 1.0
  lpip = lpips_model.forward(predict_image_lpips, gt_lpips).cpu().detach().numpy()
  lpip = np.squeeze(lpip)

  # calculate PSNR/SSIM
  gt_image = gt_image.cpu().numpy()
  ssim = structural_similarity(rendered_image, gt_image, win_size=11, multichannel=True, gaussian_weights=True)
  psnr = peak_signal_noise_ratio(rendered_image, gt_image, data_range=1.0)

  return psnr, ssim, lpip

def get_PCA_direction(data): # [..., dimension]
  M = np.mean(data, axis=0)
  normalized_data = data - M
  V = np.cov(normalized_data.T)
  _, eig_vectors = np.linalg.eig(V)
  return eig_vectors
